from pytorch_lightning import Trainer
from argparse import ArgumentParser
import torch
from tasks.ClassSummary.models.DeepCom.model import DeepCom
from tasks.ClassSummary.models.DeepCom.data import DeepComDataModule
from config.Config import BASE_DIR
from pytorch_lightning import seed_everything


def add_program_args(parser: ArgumentParser):
    parser.add_argument("--sbt_field", help="SBT词表文件", type=str,
                        default=str(BASE_DIR / "tasks/ClassSummary/resource/sbt-s.field"))
    parser.add_argument("--nl_field", help="自然语言词表文件", type=str,
                        default=str(BASE_DIR / "tasks/ClassSummary/resource/nl-s.field"))
    parser.add_argument("--check_point", help="check point path", default=None)
    parser.add_argument("--save_path", help="结果存储文件", default=None)
    parser.add_argument("--seed", help="random seed", default=1024)
    parser.add_argument("--gpu_num", help="GPUS", default=0)


def main(args):
    seed_everything(args.seed)
    sf = torch.load(args.sbt_field)
    nf = torch.load(args.nl_field)
    data = DeepComDataModule(
        train_path=None,
        val_path=None,
        test_path=args.test,
        sbt_field=sf,
        nl_field=nf,
        batch_size=args.batch_size,
    )
    model = DeepCom.load_from_checkpoint(args.check_point,
                                         code_vocab=sf.vocab,
                                         nl_vocab=nf.vocab,
                                         translate_path=args.save_path)
    trainer = Trainer(enable_checkpointing=False, gpus=args.gpu_num)

    trainer.test(model, datamodule=data)


if __name__ == '__main__':
    # 解析命令行参数
    arg_parser = ArgumentParser()
    add_program_args(arg_parser)
    DeepComDataModule.add_data_args(arg_parser)
    cli_args = arg_parser.parse_args()
    main(cli_args)
